import os
import random

import numpy as np
from PIL import Image
import torch
from torch.utils.data import Dataset
from torchvision.transforms import InterpolationMode
import cc3d
import torchvision
import matplotlib.pyplot as plt
import cv2
import matplotlib as mpl

mpl.use('TkAgg')
import skimage
from sklearn.model_selection import train_test_split
import torchio as tio

# 'NPZ/WCWW_data\\0005203451\\T1 Sag_1.2.840.113619.2.222.3596.375486.9092.1587436046.509.npz'

# transforms_dict = {
#     tio.RandomBlur(1): 0.2,
#     tio.RandomElasticDeformation(num_control_points=5): 0.4,
# }
# tr = tio.OneOf(transforms_dict)
num_cl = 6
# lw = 448
# ww = 384
# degree = (-10, 10)


def Trans(degree, ww, lw):
    train_transformI = torchvision.transforms.Compose([
        torchvision.transforms.Resize(lw, interpolation=torchvision.transforms.InterpolationMode.BICUBIC),
        torchvision.transforms.RandomRotation((-degree, degree),
                                              interpolation=torchvision.transforms.InterpolationMode.BICUBIC),

        torchvision.transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2),

        torchvision.transforms.CenterCrop(ww),

        torchvision.transforms.ToTensor(),

    ])
    return train_transformI


def TransR(degree, ww, lw):
    train_transformI = torchvision.transforms.Compose([
        torchvision.transforms.Resize(lw, interpolation=torchvision.transforms.InterpolationMode.NEAREST),
        torchvision.transforms.RandomRotation((-degree, degree),
                                              interpolation=torchvision.transforms.InterpolationMode.NEAREST),
        torchvision.transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2),

        torchvision.transforms.CenterCrop(ww),

        torchvision.transforms.ToTensor(),
    ])
    return train_transformI


def vali_trans(ww, lw):
    train_transformI = torchvision.transforms.Compose([
        torchvision.transforms.Resize(lw, interpolation=torchvision.transforms.InterpolationMode.BICUBIC),
        torchvision.transforms.CenterCrop(ww),
        torchvision.transforms.ToTensor(),
    ])
    return train_transformI


def vali_transR(ww, lw):
    train_transformI = torchvision.transforms.Compose([
        torchvision.transforms.Resize(lw, interpolation=torchvision.transforms.InterpolationMode.NEAREST),
        torchvision.transforms.CenterCrop(ww),
        torchvision.transforms.ToTensor(),

    ])
    return train_transformI


def cord_(mask):
    Z, X, Y = np.where(mask > 0)
    z, w, h = mask.shape

    cent_z = (np.max(Z) + np.min(Z)) / 2. / z
    cent_x = (np.max(X) + np.min(X)) / 2. / w
    cent_y = (np.max(Y) + np.min(Y)) / 2. / h

    z_n = (np.max(Z) - np.min(Z) + 1) / z / 1.
    x_n = (np.max(X) - np.min(X) + 1) / w / 1.
    y_n = (np.max(Y) - np.min(Y) + 1) / h / 1.

    return [cent_z, cent_x, cent_y, z_n, x_n, y_n]


class MyDataSet(Dataset):
    def __init__(self, data_list, size, lw, degree, transform=True):
        self.data_list = data_list
        self.transform = transform
        self.size = size
        self.lw = lw
        self.degree = degree

    def __len__(self):
        return len(self.data_list)

    def _encode(self, cord, classes):
        encoded = {
            "boxes": cord,
            "labels": classes
        }
        return encoded

    def get_label(self, tmp_subfolder):
        label_dic = {'HCC': 0, 'CYST': 1, 'FNH': 3, 'HA': 2, 'Hemangioma': 2, 'ICC': 4, 'META': 5, 'nodule': 1}
        if 'HCC' in tmp_subfolder:
            return label_dic['HCC']
        elif 'CYST' in tmp_subfolder:
            return label_dic['CYST']
        elif 'FNH' in tmp_subfolder:
            return label_dic['FNH']
        elif 'HA' in tmp_subfolder:
            return label_dic['HA']
        elif 'Hemangioma' in tmp_subfolder:
            return label_dic['Hemangioma']
        elif 'ICC' in tmp_subfolder:
            return label_dic['ICC']
        elif 'META' in tmp_subfolder:
            return label_dic['META']
        elif 'nodule' in tmp_subfolder:
            return label_dic['nodule']

    def mean_std(self, data):
        std_ = [0.229, 0.224, 0.225]
        mean_ = [0.485, 0.456, 0.406]
        for i in range(data.shape[0]):
            if i % 3 == 0:
                data[i, :, :] = (data[i, :, :] - mean_[0]) / std_[0]
            elif i % 3 == 1:
                data[i, :, :] = (data[i, :, :] - mean_[1]) / std_[1]
            else:
                data[i, :, :] = (data[i, :, :] - mean_[2]) / std_[2]

        return data

    def _get_data(self, data, id, seed):
        lw = self.lw
        degree = self.degree
        lyers = self.size[0]
        width = self.size[1]
        height = self.size[2]

        if self.transform:
            transformI = Trans(degree, width, lw)
            transformR = TransR(degree, width, lw)
        else:
            transformI = vali_trans(ww=width, lw=lw)
            transformR = vali_transR(ww=width, lw=lw)

        image = data['image']
        mask = data['mask']
        mask = mask>0.1
        mask = (mask * 255).astype(np.uint8)

        index = np.sum(np.sum(mask, -1), -1)
        index = np.where(index > 0)[0]
        if len(index)==0:
            print('qqqqqwwwwjdfifhfkdndbf', id)
        if len(index) < 4:
            z_, x_, y_ = np.where(mask > 0)
            start_z = max(0, min(z_) - 2)
            end_z = min(max(z_) + 2, mask.shape[0])
            start_x = max(0, min(x_) - 5)
            end_x = min(max(x_) + 5, mask.shape[1])
            start_y = max(0, min(y_) - 5)
            end_y = min(max(y_) + 5, mask.shape[2])

            new_mask = np.zeros(mask.shape).astype(np.uint8)
            new_mask[start_z:end_z, start_x:end_x, start_y:end_y] = 255
            mask = new_mask
        torch.manual_seed(seed)

        transformed_data = transformI(Image.fromarray(image[0, :, :]))
        torch.manual_seed(seed)
        transformed_mask = transformR(Image.fromarray(mask[0, :, :]))

        for i in range(image.shape[0] - 1):
            torch.manual_seed(seed)
            trans_data = transformI(Image.fromarray(image[i + 1, :, :]))
            torch.manual_seed(seed)
            trans_mask = transformR(Image.fromarray(mask[i + 1, :, :]))
            transformed_data = torch.cat([transformed_data, trans_data], dim=0)
            transformed_mask = torch.cat([transformed_mask, trans_mask], dim=0)

        transformed_data = (transformed_data.numpy() * 255)
        transformed_mask = transformed_mask.numpy()
        transformed_mask = transformed_mask > 0.1
        transformed_mask = (transformed_mask * 255).astype(np.uint8)

        transformed_data = skimage.transform.resize(transformed_data, (lyers, width, height), anti_aliasing=False,
                                                    order=1)
        transformed_mask = skimage.transform.resize(transformed_mask, (lyers, width, height), anti_aliasing=False,
                                                    order=0)

        transformed_data = transformed_data.astype(np.uint8)

        transformed_data = transformed_data / 255.
        # transformed_data = (transformed_data-np.mean(transformed_data))/ np.std(transformed_data)
        transformed_data = self.mean_std(transformed_data)
        transformed_mask = transformed_mask > 0.1
        transformed_mask = transformed_mask.astype(np.uint8)
        index = np.sum(np.sum(transformed_mask, -1), -1)
        index = np.where(index > 0)[0]
        if len(index) < 4 and len(index) != 0:
            z_, x_, y_ = np.where(transformed_mask > 0)
            start_z = max(0, min(z_) - 2)
            end_z = min(max(z_) + 2, transformed_mask.shape[0])
            start_x = max(0, min(x_) - 2)
            end_x = min(max(x_) + 2, transformed_mask.shape[1])
            start_y = max(0, min(y_) - 2)
            end_y = min(max(y_) + 2, transformed_mask.shape[2])

            new_mask = np.zeros(transformed_mask.shape).astype(np.uint8)
            new_mask[start_z:end_z, start_x:end_x, start_y:end_y] = 1
            transformed_mask = new_mask

        label_out = cc3d.connected_components(transformed_mask, connectivity=18)
        num_roi = np.max(label_out)

        cord = []
        labels = []

        max_roi_num = 5

        if num_roi == 0:
            cord.append(np.array([0.5, 0.5, 0.5, 0.1, 0.1, 0.1]))
            labels.append(num_cl)
        else:
            for j in range(num_roi):
                tmp_mask = label_out == j + 1
                tmp_cord = cord_(tmp_mask)
                cord.append(tmp_cord)
                tmp_label = self.get_label(id)
                labels.append(tmp_label)

        cord_zeros = np.zeros((max_roi_num, 6))
        labels_zeros = np.ones(max_roi_num) * num_cl
        cord = np.array(cord)
        labels = np.array(labels)

        if cord.shape[0] <= max_roi_num:
            cord_zeros[0:cord.shape[0], :] = cord
            labels_zeros[0:cord.shape[0]] = labels
        else:
            cord_zeros = cord[0:max_roi_num, :]
            labels_zeros = labels[0:max_roi_num]

        cord = cord_zeros
        labels = labels_zeros
        transformed_data = torch.from_numpy(transformed_data).float()
        cord = cord.astype(np.float32)
        cord = torch.from_numpy(cord)

        labels = torch.from_numpy(labels)
        labels = labels.long()

        transformed_mask = torch.from_numpy(transformed_mask)

        targets_list = [cord, labels]

        return [transformed_data, targets_list, transformed_mask]

    def __getitem__(self, index):
        id = self.data_list[index]
        img_path2 = os.path.join(id, 'a.npz')
        img_path3 = os.path.join(id, 'p.npz')
        img_path4 = os.path.join(id, 'v.npz')
        data2 = np.load(img_path2)
        data3 = np.load(img_path3)
        data4 = np.load(img_path4)

        seed = np.random.randint(0, 2 ** 16)

        datalist2 = self._get_data(data2, id, seed)
        datalist3 = self._get_data(data3, id, seed)
        datalist4 = self._get_data(data4, id, seed)

        return datalist2, datalist3, datalist4, id


def list_z(cord, index_z):
    cord_xy = []
    if cord.shape[0] > 0:
        for i in range(cord.shape[0]):
            tmp_z = cord[i, 0:2]
            b = np.arange(tmp_z[1] - tmp_z[0] + 1) + tmp_z[0]
            b = list(b)
            if index_z in b:
                cord_xy.append(cord[i, 2:])
    return cord_xy


def cord_trans(size, cord):
    if cord.shape[0] == 0:
        cord = torch.tensor([0, 0, 0, 0, 0, 0]).to(cord.device)
        cord = torch.unsqueeze(cord, 0)

    z = size[0]
    w = size[1]
    h = size[2]
    start_z = torch.round(cord[:, 0] * z - cord[:, 3] * z / 2 + 0.5)
    end_z = torch.round(cord[:, 0] * z + cord[:, 3] * z / 2 - 0.5)
    end_z[end_z >= z] = z - 1
    start_x = torch.round(cord[:, 1] * w - cord[:, 4] * w / 2 + 0.5)
    end_x = torch.round(cord[:, 1] * w + cord[:, 4] * w / 2 - 0.5)

    start_y = torch.round(cord[:, 2] * h - cord[:, 5] * h / 2 + 0.5)
    end_y = torch.round(cord[:, 2] * h + cord[:, 5] * h / 2 - 0.5)

    return torch.stack((start_z, end_z, start_x, end_x, start_y, end_y), -1)


if __name__ == '__main__':
    path = 'data/'
    lw = 448
    size = (48, 384, 384)
    degree = 45
    label_dic = {'HCC': 0, 'CYST': 1, 'FNH': 2, 'HA': 3, 'Hemangioma': 4, 'ICC': 5, 'META': 6, 'nodule': 7}

    folder1s = os.listdir(path)
    # folder1s.remove('.DS_Store')
    all_num = 0
    data_num = 0
    data_address = []
    labels = []
    for name in folder1s:
        tmp_folder = os.path.join(path, name)
        sub_names = os.listdir(tmp_folder)
        if '.DS_Store' in sub_names:
            sub_names.remove('.DS_Store')
        for sub_name in sub_names:
            data_num += 1
            tmp_subfolder = os.path.join(tmp_folder, sub_name)
            data_address.append(tmp_subfolder)
            if 'HCC' in tmp_subfolder:
                labels.append(label_dic['HCC'])
            elif 'CYST' in tmp_subfolder:
                labels.append(label_dic['CYST'])
            elif 'FNH' in tmp_subfolder:
                labels.append(label_dic['FNH'])
            elif 'HA' in tmp_subfolder:
                labels.append(label_dic['HA'])
            elif 'Hemangioma' in tmp_subfolder:
                labels.append(label_dic['Hemangioma'])
            elif 'ICC' in tmp_subfolder:
                labels.append(label_dic['ICC'])
            elif 'META' in tmp_subfolder:
                labels.append(label_dic['META'])
            elif 'nodule' in tmp_subfolder:
                labels.append(label_dic['nodule'])
            else:
                print(tmp_subfolder)

    X_train, X_test, y_train, y_test = train_test_split(data_address, labels, test_size=0.2, random_state=42)

    test_dataset = MyDataSet(X_train, size, lw=lw, degree=degree, transform=True)

    data_list2, data_list3, data_list4, ids = test_dataset.__getitem__(0)
    transformed_data2 = data_list2[0].numpy()
    transformed_data3 = data_list3[0].numpy()
    cord2 = torch.unsqueeze(data_list2[1][0][0, :], 0)
    cord2 = cord_trans(size=size, cord=cord2).int().numpy()
    print(cord2)
    label2 = data_list2[1][1]
    # transformed_data = np.squeeze(transformed_data)
    transformed_mask2 = data_list2[2].numpy()
    transformed_mask3 = data_list3[2].numpy()
    cord3 = torch.unsqueeze(data_list3[1][0][0, :], 0)
    cord3 = cord_trans(size, cord3).int().numpy()
    label3 = data_list3[1][1]
    # transformed_mask = np.squeeze(transformed_mask)
    Z = np.sum(np.sum(transformed_mask2, -1), -1)
    index = np.where(Z > 0)[0]
    for j in range(index.shape[0]):
        cord_xy_label = list_z(cord2, index[j])
        tmp_mask = np.zeros((size[1], size[2])).astype(np.uint(8))
        for xy in cord_xy_label:
            tmp_mask[xy[0]:xy[1], xy[2]:xy[3]] = 1
        plt.figure()
        plt.imshow(transformed_data2[index[j], :, :], cmap='gray')
        plt.contour(transformed_mask2[index[j], :, :])
        plt.contour(tmp_mask)
        plt.show()

    # Z = np.sum(np.sum(transformed_mask2, -1), -1)
    # index = np.where(Z > 0)[0]
    # for j in range(index.shape[0]):
    #
    #     plt.figure()
    #     plt.imshow(transformed_data2[index[j], :, :], cmap='gray')
    #     plt.contour(transformed_mask2[index[j], :, :])
    #     plt.show()

